#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/1/31 17:03
# @Author  : Aries
# @Site    : 
# @File    : image_tools.py
# @Software: PyCharm Community Edition
import gdal,osr


def readDEM(demfile):
    demds = gdal.Open(demfile)
    demband = demds.GetRasterBand(1)
    dem = demband.ReadAsArray()
    return dem

def readRadianTIFF(imagefile):
    imgds = gdal.Open(imagefile)
    imggt = imgds.GetGeoTransform()
    #upper left
    print('raster geotransform coeffs:', imggt[0], imggt[1], imggt[2], imggt[3], imggt[4], imggt[5])
    width = imgds.RasterXSize  # 5029
    height = imgds.RasterYSize  # 5177
    #rasterOrigin = (minx, maxy)
    band2 = imgds.GetRasterBand(2)
    b2 = band2.ReadAsArray()
    band3 = imgds.GetRasterBand(3)
    b3 = band3.ReadAsArray()
    band4 = imgds.GetRasterBand(4)
    b4 = band4.ReadAsArray()
    band5 = imgds.GetRasterBand(5)
    b5 = band5.ReadAsArray()
    band6 = imgds.GetRasterBand(6)
    b6 = band6.ReadAsArray()
    band7 = imgds.GetRasterBand(7)
    b7 = band7.ReadAsArray()
    return imggt, b2, b3, b4, b5, b6, b7


def img2array(img):
    """
    read dems to array by gdal
    :param imgfn path of geotiff
    :return narray of geotiff
    """
    img_data = gdal.Open(img)
    img_array = img_data.ReadAsArray()
    return img_array

def array2raster(newRasterfn, panTransform, array):
    cols = array.shape[1]
    rows = array.shape[0]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((panTransform[0], panTransform[1], panTransform[2], panTransform[3],
                               panTransform[4], panTransform[5]))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromEPSG(32643)
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()

def array2ascii(ncols, nrows, xllcorner, yllcorner, cellsize):
    """
    array to esri ascii raster
    :param ncols:
    :param nrows:
    :param xllcorner:
    :param yllcorner:
    :param cellsize:
    :return:
    """
    pass